import java.util.Arrays;
import java.util.PriorityQueue;

/**
 * 480. 滑动窗口中位数
 * https://leetcode-cn.com/problems/sliding-window-median/
 */
public class Solutions_480 {
    public static void main(String[] args) {
//        int[] nums = {1, 3, -1, -3, 5, 3, 6, 7};
//        int k = 3;  // output: {1, -1, -1, 3, 5, 6}

//        int k = 1;  // output: {1, 3, -1, -3, 5, 3, 6, 7}
//        int k = 4;  // output: {0, 1, 1, 4.0, 5.5}

//        int[] nums = {1, 4, 2, 3};
//        int k = 4;  // output: {2.5}

        int[] nums = {2147483647,2147483647};
        int k = 2;  // output: {2147483647.0}

        double[] result = medianSlidingWindow(nums, k);
        System.out.println(Arrays.toString(result));
    }

    // 小顶堆（小值在堆顶，堆中存储是的大于等于中位数的数）
    private static PriorityQueue<Integer> minHeap = null;
    // 大顶堆（大值在堆顶，堆中存储是的小于等于中位数的数）
    private static PriorityQueue<Integer> maxHeap = null;

    /**
     * 解法一：堆（76ms）
     */
    public static double[] medianSlidingWindow(int[] nums, int k) {
        minHeap = new PriorityQueue<>();
        // 数值溢出
//        maxHeap = new PriorityQueue<>((o1, o2) -> o2 - o1);
        maxHeap = new PriorityQueue<>((o1, o2) -> Integer.compare(o2, o1));
        int len = nums.length;
        double[] res = new double[len - k + 1];
        for (int i = 0; i < len; i++) {
            addNum(nums[i]);
            if (i >= k - 1) {
                double midNum = getMidNum();
                res[i - k + 1] = midNum;
                removeNum(nums[i + 1 - k], midNum);
            }
        }
        return res;
    }

    /**
     * 添加元素到堆中
     */
    public static void addNum(int num) {
        if (minHeap.size() <= maxHeap.size()) {
            // 放大堆，然后大堆堆顶元素移动到小堆
            maxHeap.add(num);
            minHeap.add(maxHeap.poll());
        } else {
            // 放小堆，然后小堆堆顶元素移动到大堆
            minHeap.add(num);
            maxHeap.add(minHeap.poll());
        }
    }

    /**
     * 获取中位数
     */
    public static double getMidNum() {
        if (minHeap.size() > maxHeap.size()) {
            return minHeap.peek();
        } else if (maxHeap.size() > minHeap.size()) {
            return maxHeap.peek();
        } else {
            int p1 = minHeap.peek();
            int p2 = maxHeap.peek();
            // 防止溢出
//            return (p1 + p2) / 2.0;
//            return ((long) p1 + p2) / 2.0;
            return (p1 / 2.0) + (p2 / 2.0);
        }
    }

    /**
     * 移除 midNum
     * 若 midNum 大于等于中位数，那么一定在 minHeap 中，否则在 maxHeap 中
     * 直接移除元素即可，因为后面 add 时，会根据堆大小来移动元素
     */
    public static void removeNum(int num, double midNum) {
        if (num >= midNum) {
            minHeap.remove(num);
        } else {
            maxHeap.remove(num);
        }
    }
}
